import copy
import time
from datetime import datetime
import random
from typing import Any, Dict, Optional
import numpy as np
import gymnasium as gym
from gymnasium import spaces
import cv2
import platform

from gym_pih.hardware.RealsenseCamera import Camera as RealsenseCamera
from gym_pih.hardware.ft300 import RobotiqFt300 as FT300
from gym_pih.hardware.pid import PIDForceController
from gym_pih.hardware.calculate_slope import calculate_slope
from gym_pih.hardware.LightSource import LightSource
from gym_pih.reward_classifier.load_classifier_for_v3 import Classifier
import rtde_control
import rtde_receive

if platform.system() == 'Windows':
    from gym_pih.hardware.Imagingsource_for_windows import Camera as ImagingsourceCamera
elif platform.system() == 'Linux':
    # from gym_pih.hardware.Imagingsource_for_linux import Camera as ImagingsourceCamera
    from gym_pih.hardware.asyn_camera import Camera as ImagingsourceCamera
else:
    raise Exception("Unsupported platform")

class TypecEnv(gym.Env):
    def __init__(self,
                 target_pos: np.ndarray = np.array(
                     [0.6874792884908796,
                     -0.1290049028269131,
                     0.17394818223574388,
                     0.00014219014866164234,
                     -3.1375895073275664,
                     2.166236079437557e-05] ),
                 reward_shape="sparse",
                 train_env='normal_env',
                 fake_env=False,
                 curriculum=False,
                 save_video = False,
                 render = False
                 ):
        self.params_pos_error=0   #孔的定位误差范围
        self.action_range=0.001   #动作范围1mm
        self.initial_error=0.01   #初始位置与孔的最大距离
        self.ministep_num=6       #每次动作分8步执行
        
        self.train_env=train_env
        self.ip="192.168.3.140"
        self.velocity = 0.1
        self.acceleration = 0.05
        self.dt = 1.0/125  # 8ms
        self.lookahead_time = 0.1
        self.gain = 300
        
        assert reward_shape in ["sparse","dense"],f"reward_shape must be in ['sparse','dense'],but got {reward_shape}"
        self.reward_shape=reward_shape
        self.action_space = spaces.Box(low=-1, high=1, shape=(3,), dtype=np.float32)
        if train_env=='serl':
            self.observation_space = gym.spaces.Dict(
                {
                    "state": gym.spaces.Dict(
                        {
                            "position": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(3,), dtype=np.float32),
                            "force":gym.spaces.Box(low=-np.inf, high=np.inf, shape=(3,), dtype=np.float32),
                            "slope":gym.spaces.Box(low=-np.inf, high=np.inf, shape=(1,), dtype=np.float32),
                        }
                    ),
                    "images": gym.spaces.Dict(
                        {
                            "wrist": gym.spaces.Box(
                                low=0,
                                high=255,
                                shape=(64, 64, 3),
                                dtype=np.uint8,
                            ),
                            "side": gym.spaces.Box(
                                low=0,
                                high=255,
                                shape=(64, 64, 3),
                                dtype=np.uint8,
                            ),
                        }
                    ),
                }
            )

        else:
            self.observation_space = gym.spaces.Dict(
                {"state":gym.spaces.Box(low=-np.inf, high=np.inf, shape=(7,), dtype=np.float32),
                 "wrist":gym.spaces.Box(low=0,high=1,shape=(3,64,64),dtype=np.float32),
                 "side":gym.spaces.Box(low=0,high=1,shape=(3,64,64),dtype=np.float32),}
            )

        

        self.pre_distance=None
        self.pre_slope=0
        self.curriculum = curriculum
        self.curriculum_difficulty=0
        self.success_rate = 0
        # self.switch_demo = 0
        assert target_pos.shape==(6,),f"target_pos must be a 6-dim ndarray,but got {target_pos.shape}"
        self.target_pos=target_pos.copy() 
        self.pos_error=np.random.uniform(-self.params_pos_error,self.params_pos_error,2)
        self.target_pos[:2]+=self.pos_error#孔的位置有定位误差
        self.pose=np.zeros(6)#用于控制机器人运动
        self.boundary=np.zeros((2,3))#用于限制机器人的运动范围
        self.boundary[0]=self.target_pos[:3]-0.025
        self.boundary[0][2]=self.target_pos[2]-0.003
        self.boundary[1]=self.target_pos[:3]+0.025
        self.calculate_boundary=self.boundary.copy()

        if fake_env:
            return
        self.Light=LightSource()
        self.Light.Light_Open(6)
        self.Realsense_cam = RealsenseCamera()
        self.Realsense_cam.Open()
        self.Imagingsource_cam = ImagingsourceCamera()
        self.Imagingsource_cam.Open()
        if platform.system() == 'Windows':
            self.Imagingsource_cam.SetProperty("Gain","Value",6.9)
            self.Imagingsource_cam.SetProperty("Exposure","Value",float(1/60))
        self.ft=FT300()
        # self.gpio=GPIO()
        # self.gpio.IO_InitPin(0,0,2)
        self.action_controller=PIDForceController(kp=self.action_range/self.ministep_num/4,
                                                  ki=self.action_range/self.ministep_num/100/150,
                                                  kd=self.action_range/self.ministep_num/4)
        self.rtde_c = rtde_control.RTDEControlInterface(self.ip)
        self.rtde_r = rtde_receive.RTDEReceiveInterface(self.ip)

        self.classifier=Classifier()
        self.save_video=save_video
        self.render=render
        if self.save_video:
            self.recording_frames = []

    def _get_obs(self):
        target_pos=self.target_pos.copy()[:3]
        pos_ee=self.rtde_r.getActualTCPPose()[:3]
        force=self.ft.forceTorque[:3]
        _,wrist_image=self.Realsense_cam.GetImage(isShow=False)
        _,side_image=self.Imagingsource_cam.GetImage()
        distance=pos_ee-target_pos
        pos_ee=np.array(pos_ee,dtype=np.float32)
        distance=np.array(distance,dtype=np.float32)
        force=np.array(force,dtype=np.float32)
        wrist_image=wrist_image[150:214,215:279]
        done,slope=calculate_slope(side_image[130:,250:1020])
        if not done:
            slope=self.pre_slope
        else:
            self.pre_slope=slope
        slope=np.array([abs(slope)],dtype=np.float32)
        side_image=side_image[60:700, 260:900]
        side_image=cv2.resize(side_image,(64,64))
        if self.render:
            cv2.namedWindow("wrist",cv2.WINDOW_NORMAL)
            cv2.namedWindow("side",cv2.WINDOW_NORMAL)
            cv2.resizeWindow("wrist",320,320)
            cv2.resizeWindow("side",320,320)
            cv2.imshow("wrist",wrist_image)
            cv2.imshow("side",side_image)
            cv2.waitKey(1)
        if self.save_video:
            self.recording_frames.append(np.concatenate((wrist_image,side_image),axis=0))
        if self.train_env=='serl':
            observation={"state":{"position":pos_ee,"force":force,"slope":slope},
                         "images":{"wrist":wrist_image,"side":side_image}}
            observation=self._normalize(observation)
        else:
            state=np.concatenate((pos_ee,force,slope))
            wrist_image= np.transpose(wrist_image, (2, 0, 1))
            side_image = np.transpose(side_image, (2, 0, 1))
            observation={"state":state,"wrist":wrist_image,"side":side_image}
            observation=self._normalize(observation)
        return observation
    
    def _get_info(self):
        return {"success_rate":self.success_rate,
                "curriculum_difficulty":self.curriculum_difficulty,
                "switch_demo":0}
    
    def _get_reward(self,observation,terminated):
        if self.train_env=="serl":
            reward=-0.005
            insert_state=self.classifier.predict(observation)
            reward+=insert_state*0.002
            if insert_state:
                slope=(observation["state"]["slope"]).item()
                reward+=(slope-0.5)*0.002
            if terminated:
                reward=1
        else:
            reward = -0.005
            obs={"images":{"wrist":np.transpose(observation["wrist"],(1,2,0)),"side":np.transpose(observation["side"],(1,2,0))}}
            insert_state = self.classifier.predict(obs)
            reward += insert_state * 0.002
            if insert_state:
                slope = (observation["state"][6]).item()
                reward += (slope - 0.5) * 0.002
            if terminated:
                reward = 1
        return reward
    
    def _distance(self,distance_array):
        return np.linalg.norm(distance_array)
    
    def reset(self, seed: int = None, options: Optional[Dict[str, Any]] = None):
        """先向上移动拔出插头，再复位"""
        if self.save_video:
            time.sleep(0.5)
            _,wrist_image=self.Realsense_cam.GetImage(isShow=False)
            _,side_image=self.Imagingsource_cam.GetImage()
            wrist_image=wrist_image[150:214,215:279]
            side_image=side_image[60:700, 260:900]
            side_image=cv2.resize(side_image,(64,64))
            self.recording_frames.append(np.concatenate((wrist_image,side_image),axis=0))
            self.save_video_recording()
        self.rtde_c.servoStop()        #ur进入servo模式后，必须先停止才能再次用move
        self._go_to_reset_pose()

        self.pose=self.rtde_r.getActualTCPPose()
        self.pre_distance=None
        observation=self._get_obs()
        info=self._get_info()
        self.switch_demo = 0
        if not self.curriculum:
            self.curriculum_difficulty=random.uniform(0.1,1.0)
        return observation,info
    
    def step(self, action):
        self._miniServoL(action*self.action_range,self.ministep_num)
        terminated=self.rtde_r.getDigitalInState(0)
        observation=self._get_obs()
        reward=self._get_reward(observation,terminated)
        info=self._get_info()

        return observation, reward, terminated, False, info
    
    # def _go_to_reset_pose(self):
    #     init_pose=self.target_pos.copy()
    #     init_pose[:2] += np.random.uniform(-self.initial_error,self.initial_error, 2)
    #     init_pose[2]+=0.015
    #     transition_pose=self.rtde_r.getActualTCPPose()
    #     if self.curriculum_difficulty==0:
    #         transition_pose[2]=init_pose[2]
    #         self.rtde_c.moveL(transition_pose, 0.2, 0.05)
    #         self.rtde_c.moveL(init_pose, 0.2, 0.05)
    #     elif self.curriculum_difficulty>=0.1 and self.curriculum_difficulty<=0.6:
    #         transition_pose[1]+=self.randomvariate(-self.curriculum_difficulty*0.04/0.6,0) if init_pose[1]<self.target_pos[1] \
    #             else self.randomvariate(0,self.curriculum_difficulty*0.04/0.6)
    #         self.rtde_c.moveL(transition_pose, 0.2, 0.05)
    #         self.rtde_c.moveL(init_pose, 0.2, 0.05)
    #     else:
    #         transition_pose[1]+=0.04 if init_pose[1]<self.target_pos[1] else -0.04
    #         self.rtde_c.moveL(transition_pose, 0.2, 0.05)
    #         transition_pose[2]+=0.015
    #         self.rtde_c.moveL(transition_pose, 0.2, 0.05)
    #         transition_pose[:2]=init_pose[:2]
    #         self.rtde_c.moveL(transition_pose, 0.2, 0.05)
    #         init_pose[2]-=self.randomvariate(0.005,0.012*self.curriculum_difficulty)
    #         self.rtde_c.moveL(init_pose, 0.2, 0.05)

    def _go_to_reset_pose(self):
        init_pose=self.target_pos.copy()
        init_pose[:2] += np.random.uniform(-self.initial_error,self.initial_error, 2)
        init_pose[2]+=0.015
        transition_pose=self.rtde_r.getActualTCPPose()
        transition_pose[2]=init_pose[2]
        self.rtde_c.moveL(transition_pose, 0.2, 0.05)
        if self.curriculum_difficulty==0:
            self.rtde_c.moveL(init_pose, 0.2, 0.05)
        elif self.curriculum_difficulty>=0.1:
            random_bit = random.randint(0, 1)
            transition_pose[0] = self.target_pos[0]
            if random_bit == 0:
                transition_pose[1] = -0.19
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                transition_pose[2] = 0.158
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                transition_pose[1] = -0.1664
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                transition_pose[1] += 0.046*self.randomvariate(0, self.curriculum_difficulty)
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                transition_pose[2] = self.target_pos[2]+0.015

            if random_bit == 1:
                transition_pose[1] = -0.055
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                transition_pose[2] = 0.158
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                transition_pose[1] = -0.073
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                transition_pose[1] -= 0.046*self.randomvariate(0, self.curriculum_difficulty)
                self.rtde_c.moveL(transition_pose, 0.2, 0.05)
                transition_pose[2] = self.target_pos[2]+0.015
            self.rtde_c.moveL(transition_pose, 0.2, 0.05)
            self.rtde_c.moveL(init_pose, 0.2, 0.05)



    def _miniServoL(self, action, step=4):
        """
        输入一次动作，分多步执行

        :param action: 动作
        :param step: 步数，默认为4
        """
        miniaction = action / step
        for i in range(step):
            force=self.ft.GetForce()
            delta=self.action_controller._calPIDForce(force)
            self.pose[:3] += miniaction+delta[:3]
            self.pose[:3]=np.clip(self.pose[:3],self.boundary[0],self.boundary[1])
            t_start=self.rtde_c.initPeriod()
            self.rtde_c.servoL(self.pose, self.velocity, self.acceleration, self.dt, self.lookahead_time, self.gain)
            self.rtde_c.waitPeriod(t_start)

    def save_video_recording(self):
        try:
            if len(self.recording_frames):
                video_writer = cv2.VideoWriter(
                    f'/home/wk/RL-sim/data/videos/env-v3/{datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}.mp4',
                    cv2.VideoWriter_fourcc(*"mp4v"),
                    10,
                    self.recording_frames[0].shape[:2][::-1],
                )
                for frame in self.recording_frames:
                    video_writer.write(frame)
                video_writer.release()
            self.recording_frames.clear()
        except Exception as e:
            print(f"Failed to save video: {e}")

    def _normalize(self,observation):
        """
        观测值归一化

        :param observation: 观测值
        :return: 归一化后的观测值
        """
        normal_obs = copy.deepcopy(observation)
        if self.train_env=='serl':
            normal_obs['state']['position']=((normal_obs['state']['position'] - self.calculate_boundary[0]) / (self.calculate_boundary[1] - self.calculate_boundary[0])).astype(np.float32)
            normal_obs['state']['force']=np.clip(normal_obs['state']['force'],-10,10)/10
            normal_obs['state']['slope']/=90
        else:
            normal_obs["state"][0:3]=((normal_obs["state"][0:3]-self.calculate_boundary[0]) / (self.calculate_boundary[1] - self.calculate_boundary[0])).astype(np.float32)
            normal_obs["state"][3:6]=np.clip(normal_obs["state"][3:6],-10,10)/10
            normal_obs["state"][6]/=90
            normal_obs["wrist"]=(normal_obs["wrist"]/255.0).astype(np.float32)
            normal_obs["side"]=(normal_obs["side"]/255.0).astype(np.float32)
        return normal_obs
    
    def randomvariate(self,xmin=0, xmax=1):
        """
        Rejection method for generating a single random number from a probability distribution.

        Parameters:
            pdf (function): Probability density function from which you want to generate random numbers.
            xmin (float): Lower bound of the desired range.
            xmax (float): Upper bound of the desired range.

        Returns:
            float: A random number following the input probability distribution.
        """

        def pdf(x):
            return abs(x)*2+0.1
        
        # Calculate the minimal and maximum values of the PDF in the desired interval.
        # The rejection method needs these values in order to work properly.
        # x = np.linspace(xmin, xmax, 1000)
        # y = pdf(x)
        # pmin = 0.
        # pmax = y.max()

        #这里的pdf是一个线性函数，所以pmin=0，pmax=pdf(xmax)
        pmin = 0
        pmax = pdf(xmax)

        # Keep generating numbers until we accept one
        while True:
            x_prime = np.random.uniform(xmin, xmax)  # x'
            y_prime = np.random.uniform(pmin, pmax)  # y'

            if y_prime < pdf(x_prime):
                return x_prime
    
    def close(self):
        self.rtde_c.servoStop()
        self.rtde_c.stopScript()
        self.rtde_c.disconnect()
        self.rtde_r.disconnect()
        self.ft.DisConnect()
        self.Light.Light_Close()
        self.Light.port_close()
        self.Imagingsource_cam.Close()
